///////////////////////////////////////////////////////////////////////
// File:        OcrSink.cpp
// Description: 
// Author:      Philip Kommnick
// Created:     Sun September 13 13:41:56 CEST 2009
//
// This program is open source. For license terms, see the LICENSE file.
//
///////////////////////////////////////////////////////////////////////

#include "StdAfx.h"
#include "OcrSink.h"
#include "resource.h"
#include "ocrlanguages.h"

typedef struct { wchar_t str[4]; unsigned int id; } LANG;

static LANG langMap[] = {
	{L"nor", IDS_OCRNOR}, {L"pol", IDS_OCRPOL}, {L"por", IDS_OCRPOR},
	{L"ron", IDS_OCRRON}, {L"rus", IDS_OCRRUS}, {L"slk", IDS_OCRSLK},
	{L"slv", IDS_OCRSLV}, {L"spa", IDS_OCRSPA}, {L"srp", IDS_OCRSRP},
	{L"swe", IDS_OCRSWE}, {L"tgl", IDS_OCRTGL}, {L"tur", IDS_OCRTUR},
	{L"ukr", IDS_OCRUKR}, {L"vie", IDS_OCRVIE}, {L"bul", IDS_OCRBUL},
	{L"cat", IDS_OCRCAT}, {L"ces", IDS_OCRCES}, {L"dan", IDS_OCRDAN},
	{L"deu", IDS_OCRDEU}, {L"ell", IDS_OCRELL}, {L"eng", IDS_OCRENG},
	{L"fin", IDS_OCRFIN}, {L"fra", IDS_OCRFRA}, {L"hun", IDS_OCRHUN},
	{L"ind", IDS_OCRIND}, {L"ita", IDS_OCRITA}, {L"lav", IDS_OCRLAV},
	{L"lit", IDS_OCRLIT}, {L"nld", IDS_OCRNLD}
};

std::vector<size_t> OcrSink::foundLanguageIDs;
std::vector<std::wstring> OcrSink::foundLanguageStrs;

OcrSink::OcrSink(void) : rememberLang(false), ocrThread(NULL), languagePage(-1)
{
	TypesList tl;
	tl.push_back(MAKEFTYPE(L"Text Document", L"*.txt"));
	InitFileDlg(tl);

	hideButton.nButtonID = IDOK;
	hideButton.pszButtonText = MAKEINTRESOURCE(IDS_OCRHIDE);

	ZeroMemory(&tdcProgress, sizeof(TASKDIALOGCONFIG));
	tdcProgress.cbSize = sizeof(TASKDIALOGCONFIG);
	tdcProgress.hwndParent = MainWin::hWnd;
	tdcProgress.hInstance = tpInstance;
	tdcProgress.dwFlags = TDF_SHOW_MARQUEE_PROGRESS_BAR | TDF_CAN_BE_MINIMIZED | TDF_CALLBACK_TIMER;
	tdcProgress.pButtons = &hideButton;
	tdcProgress.cButtons = 1;
	tdcProgress.pszWindowTitle = MAKEINTRESOURCE(IDS_APPNAME);
	tdcProgress.pszMainIcon = MAKEINTRESOURCE(IDI_APP);
	tdcProgress.pszMainInstruction = MAKEINTRESOURCE(IDS_OCRWORKMAIN);
	tdcProgress.pszContent = MAKEINTRESOURCE(IDS_OCRWORKCONT);
	tdcProgress.lpCallbackData = reinterpret_cast<LONG_PTR>(this);
	tdcProgress.pfCallback = OcrSink::ProgressCallback;

	ZeroMemory(langButtons, sizeof(langButtons));

	WIN32_FIND_DATA fileInfo;
	ZeroMemory(&fileInfo, sizeof(WIN32_FIND_DATA));
	HANDLE langFiles = FindFirstFile(L".\\tessdata\\*.traineddata", &fileInfo);
	wchar_t langBuffer[256];
	if (langFiles!=INVALID_HANDLE_VALUE) {
		do {
			int langIndex = -1;
			for (int i=0; i<_countof(langMap); i++) {
				if (!_wcsnicmp(langMap[i].str, fileInfo.cFileName, 3)) {
					langIndex = i;	
					break;
				}
			}
			if (langIndex!=-1) {
				foundLanguageIDs.push_back(langIndex);
				LoadString(tpInstance, IDS_OCRSTART+langIndex, langBuffer, sizeof(langBuffer));
				foundLanguageStrs.push_back(langBuffer);
			}
		} while(FindNextFile(langFiles, &fileInfo));
		FindClose(langFiles);
	}

	ZeroMemory(&tdcLanguage, sizeof(TASKDIALOGCONFIG));
	tdcLanguage.cbSize = sizeof(TASKDIALOGCONFIG);
	tdcLanguage.hwndParent = MainWin::hWnd;
	tdcLanguage.hInstance = tpInstance;
	tdcLanguage.dwFlags = TDF_ALLOW_DIALOG_CANCELLATION | TDF_USE_COMMAND_LINKS;
	tdcLanguage.pszWindowTitle = MAKEINTRESOURCE(IDS_APPNAME);
	tdcLanguage.pszMainIcon = MAKEINTRESOURCE(IDI_OCR);
	tdcLanguage.pszMainInstruction = MAKEINTRESOURCE(IDS_OCRINST);
	tdcLanguage.pszContent = MAKEINTRESOURCE(IDS_OCRCONT);
	tdcLanguage.pszVerificationText = MAKEINTRESOURCE(IDS_TDCREMEMBER);
	tdcLanguage.pfCallback = OcrSink::LanguageCallback;
	tdcLanguage.lpCallbackData = reinterpret_cast<LONG_PTR>(this);
	tdcLanguage.pButtons = langButtons;
	tdcLanguage.cButtons = (foundLanguageIDs.size()>itemsPerPage) ? itemsPerPage : foundLanguageIDs.size();
	LangNextPage(); // init language buttons for first use
}

OcrSink::~OcrSink(void)
{
}

void OcrSink::LangNextPage()
{
	languagePage = (languagePage+1)%((foundLanguageIDs.size()/itemsPerPage)+1);
	int maxLoop = std::min<int>(itemsPerPage-1, (foundLanguageIDs.size() - languagePage*(itemsPerPage-1)));
	for(int i=0; i<maxLoop; i++) {
		langButtons[i].nButtonID = 1000 + foundLanguageIDs[languagePage*(itemsPerPage-1)+i];
		langButtons[i].pszButtonText = foundLanguageStrs[languagePage*(itemsPerPage-1)+i].c_str();
	}
	tdcLanguage.cButtons = maxLoop;
	if (foundLanguageIDs.size()==itemsPerPage) {
		langButtons[maxLoop].nButtonID = 1000 + foundLanguageIDs[itemsPerPage-1];
		langButtons[maxLoop].pszButtonText = foundLanguageStrs[itemsPerPage-1].c_str();
		tdcLanguage.cButtons += 1;
	} else if (foundLanguageIDs.size()>itemsPerPage) {
		langButtons[maxLoop].nButtonID = 999;
		langButtons[maxLoop].pszButtonText = MAKEINTRESOURCE(IDS_OCRLANGMORE);
		tdcLanguage.cButtons += 1;
	}
}

HRESULT CALLBACK OcrSink::LanguageCallback(HWND hwnd, UINT uNotification, WPARAM wParam, LPARAM lParam, LONG_PTR dwRefData)
{
	size_t pageID = 0;
	OcrSink *that = reinterpret_cast<OcrSink*>(dwRefData);
	if (uNotification==TDN_BUTTON_CLICKED && wParam==999) {
		that->LangNextPage();
		SendMessage(hwnd, TDM_NAVIGATE_PAGE, 0, (LPARAM)&(that->tdcLanguage));
		return S_FALSE;
	}
	return S_OK;
}

HRESULT CALLBACK OcrSink::ProgressCallback(HWND hwnd, UINT uNotification, WPARAM wParam, LPARAM lParam, LONG_PTR dwRefData)
{
	OcrSink *that = reinterpret_cast<OcrSink*>(dwRefData);
	switch(uNotification) {
		case TDN_CREATED:
			SendMessage(hwnd, TDM_SET_PROGRESS_BAR_MARQUEE, TRUE, 50);
			break;
		case TDN_TIMER:
			if (WaitForSingleObject(that->ocrThread, 0)!=WAIT_TIMEOUT)
				SendMessage(hwnd, TDM_CLICK_BUTTON, IDOK, 0);
			break;
	}
	return S_OK;
}

void OcrSink::Output(Gdiplus::Bitmap &bmp)
{
	if (!RequestFileName(L"Save ocr data to... "))
		return; // user pressed cancel - do nothing
	if (!rememberLang) {
		int retButton, chckBoxState;
		switch (foundLanguageIDs.size()) {
			case 0:
				throw std::exception("no dictionary found");
				break;
			case 1:
				chckBoxState = TRUE;
				retButton = 1000+foundLanguageIDs[0];
				break;
			default:
				if (FAILED(TaskDialogIndirect ( &tdcLanguage, &retButton, NULL, &chckBoxState )) || retButton==IDCANCEL)
					throw std::exception("error opening language selection dialog");
		}
		rememberLang = (chckBoxState==TRUE);
		sprintf_s(lang, sizeof(lang), "%S", langMap[retButton-1000].str);
	}
	// prepare the image data
	Gdiplus::Rect rect(0, 0, bmp.GetWidth(), bmp.GetHeight());
	bmp.LockBits(&rect, Gdiplus::ImageLockModeRead, PixelFormat24bppRGB, &bitmapData);
	// start processing
	DWORD tID;
	ocrThread = CreateThread(NULL, 0, OcrSink::OcrFunc, this, 0, &tID);
	TaskDialogIndirect(&tdcProgress, NULL, NULL, NULL);
	WaitForSingleObject(ocrThread, INFINITE);
	ocrThread = NULL;
	bmp.UnlockBits(&bitmapData);
}

DWORD WINAPI OcrSink::OcrFunc(LPVOID lpParameter)
{
	OcrSink *that = reinterpret_cast<OcrSink*>(lpParameter);
	// initialize the ocr engine	
	tesseract::TessBaseAPI ocrApi;
	if (ocrApi.Init("tessdata", that->lang)!=0)
		throw std::exception("error initializing tesseract");
	// ocr it
	ocrApi.SetImage(reinterpret_cast<unsigned char*>(that->bitmapData.Scan0), that->bitmapData.Width, that->bitmapData.Height, 3, that->bitmapData.Stride);
	char *ocrText = NULL;
	ocrText = ocrApi.GetUTF8Text();

	// write to file
	std::ofstream outFile(that->GetFileName().c_str(), std::ios::out | std::ios::trunc);
	outFile << ocrText;
	outFile.close();
	// cleanup
	ocrApi.End();
	ocrApi.DeleteUTF8Text(ocrText);
	//delete[] ocrText; // this will trigger a memory manager assertion, but can be ignored
	CloseHandle(that->ocrThread);
	return 0;
}
